#include "wrapper_sstructmg.h"

IterativeSolver<idx_t, data_t, data_t, data_t> * solver  = nullptr;
Solver         <idx_t, data_t, data_t>         * precond = nullptr;

double  min_times[NUM_KRYLOV_RECORD], 
        max_times[NUM_KRYLOV_RECORD],
        avg_times[NUM_KRYLOV_RECORD];

void buildup_solver(std::string its_name, std::string prc_name, const TEST_CONFIG & config_file)
{
    assert(solver == nullptr);
    assert(precond== nullptr);
    
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    const int my_nblks = config_file.my_nblks;
    const int restart = config_file.restart_len;
    // 设置求解器
    if (its_name == "GCR") {
        using GCR = GCRSolver<idx_t, data_t, data_t, data_t>;
        solver = new GCR();
        ((GCR*)solver)->SetInnerIterMax(restart);
    } else if (its_name == "CG") {
        solver = new CGSolver<idx_t, data_t, data_t, data_t>;
    } else if (its_name == "GMRES") {
        using GMRES = GMRESSolver<idx_t, data_t, data_t, data_t>;
        solver = new GMRES();
        ((GMRES*)solver)->SetRestartlen(restart);
    } else {
        if (my_pid == 0) printf("INVALID iterative solver name of %s\nOnly GCR, CG, GMRES, FGMRES available\n", its_name.c_str());
        MPI_Abort(MPI_COMM_WORLD, -1);
    }
    solver->SetMaxIter(config_file.max_iter);
    solver->SetRelTol(config_file.rtol);
    solver->SetAbsTol(config_file.atol);

    // 设置预条件子
    /*
    if (strstr(prc_name.c_str(), "PJ")) {
        if (my_pid == 0) printf("  using \033[1;35mpointwise-Jacobi\033[0m as preconditioner\n");
        precond = new PointJacobi<idx_t, data_t, data_t, data_t>();
    } else if (strstr(prc_name.c_str(), "PGS")) {
        using PGS = PointGS<idx_t, data_t, data_t, data_t>;
        PGS::SCAN_TYPE type = PGS::SYMMETRIC;				
        if      (strstr(prc_name.c_str(), "F")) type = PGS::FORWARD;
        else if (strstr(prc_name.c_str(), "B")) type = PGS::BACKWARD;
        else if (strstr(prc_name.c_str(), "S")) type = PGS::SYMMETRIC;
        if (my_pid == 0) printf("  using \033[1;35mpointwise-GS %d\033[0m as preconditioner\n", type);
        precond = new PGS(type);
    } else if (strstr(prc_name.c_str(), "LJ")) {
        if (my_pid == 0) printf("  using \033[1;35mlinewise-Jacobi\033[0m as preconditioner\n");
        idx_t dir = 2;
        if 		(strstr(prc_name.c_str(), "-0")) dir = 0;
        else if (strstr(prc_name.c_str(), "-1")) dir = 1;
        else if (strstr(prc_name.c_str(), "-2")) dir = 2;
        idx_t line_dirs [my_nblks];
        for (idx_t k = 0; k < my_nblks; k++) line_dirs[k] = dir;
        precond = new LineJacobi<idx_t, data_t, data_t, data_t>(my_nblks, line_dirs);
    } else if (prc_name == "ADI") {
        if (my_pid == 0) printf("  using \033[1;35mADI\033[0m as preconditioner\n");
        idx_t first_dirs[my_nblks], second_dirs[my_nblks];
        for (idx_t k = 0; k < my_nblks; k++) {
            first_dirs [k] = 1;
            second_dirs[k] = 2;
        }
        precond = new AlterDirectIter<idx_t, data_t, data_t, data_t>(my_nblks, first_dirs, second_dirs);
    } else if (strstr(prc_name.c_str(), "BILU")) {
        __int128_t ILU_mask = stencil_mask_3d7;
        if 	    (strstr(prc_name.c_str(), "2d5i1i2")) ILU_mask = stencil_mask_2d5i1i2;
        else if (strstr(prc_name.c_str(), "2d5i0i1")) ILU_mask = stencil_mask_2d5i0i1;
        else if (strstr(prc_name.c_str(), "2d9i1i2")) ILU_mask = stencil_mask_2d9i1i2;
        else if (strstr(prc_name.c_str(), "2d9i0i1")) ILU_mask = stencil_mask_2d9i0i1;
        else if (strstr(prc_name.c_str(), "3d7" )) ILU_mask = stencil_mask_3d7 ;
        else if (strstr(prc_name.c_str(), "3d19")) ILU_mask = stencil_mask_3d19;
        else if (strstr(prc_name.c_str(), "3d27")) ILU_mask = stencil_mask_3d27;
        else {
            if (my_pid == 0) {
                char buf[130];
                printf("Error: unsupported mask of Box-ILU %s\n", print_int128(ILU_mask, buf));
            }
            MPI_Abort(MPI_COMM_WORLD, -20230121);
        }
        if (my_pid == 0) {
            char buf[130];
            printf("  using \033[1;35mBox-ILU mask %s\033[0m as preconditioner\n", print_int128(ILU_mask, buf));
        }
        __int128_t ilu_masks [my_nblks];
        for (idx_t k = 0; k < my_nblks; k++) ilu_masks[k] = ILU_mask;
        precond = new BoxILU<idx_t, data_t, data_t, data_t>(my_nblks, ilu_masks);
    } else if (prc_name == "BoxLU") {
        if (my_pid == 0) printf("  using \033[1;35mBox-LU (boxwise direct solver)\033[0m as preconditioner\n");
        precond = new BoxLU<idx_t, data_t, data_t, data_t>();
    } else if (prc_name == "BlkLU") {
        if (my_pid == 0) printf("  using \033[1;35mBlock-LU (blockwise direct solver)\033[0m as preconditioner\n");
        precond = new BlockLU<idx_t, data_t, data_t, data_t>();
    } else if (prc_name == "GlbLU") {
        if (my_pid == 0) printf("  using \033[1;35mGlobal-LU (global direct solver)\033[0m as preconditioner\n");
        precond = new GlobalLU<idx_t, data_t, data_t, data_t>();
    } else */
    if (prc_name == "GMG") {
        using MG = MultiGrid<idx_t, data_t, data_t, data_t>;
        precond = new MG(config_file.config_mg_file);
        ((MG*)precond)->ParseConfig(*(config_file.ssgrid_ptr));
    }
    else {
        if (my_pid == 0) printf("No Precond!\n");
    }

    if (precond != nullptr) solver->SetPreconditioner(*precond);
}

void setup_and_solve(const par_SstructMatrix<idx_t, data_t, data_t> & ss_A,
    const par_SstructVector<idx_t, data_t> & b, par_SstructVector<idx_t, data_t> & x,
    TEST_RECORD & record)
{
    MPI_Barrier(MPI_COMM_WORLD);
    double t0 = wall_time();
    solver->SetOperator(ss_A);
    t0 = wall_time() - t0;

    MPI_Barrier(MPI_COMM_WORLD);
    double t1 = wall_time();
    solver->Mult(b, x, false);
    t1 = wall_time() - t1;

    record.setup = (precond != nullptr) ? precond->setup_time : 0.0;
    record.solve = t1;
    record.iter  = solver->GetNumIterations();
}

void destroy_solver()
{
    // stat_sstructmg_pattern();
    if (precond) { delete precond; precond = nullptr; }
    if (solver ) { delete solver ; solver  = nullptr; }
}

void check_residual(const par_SstructMatrix<idx_t, data_t, data_t> & ss_A,
    const par_SstructVector<idx_t, data_t> & ss_x, const par_SstructVector<idx_t, data_t> & ss_b,
    par_SstructVector<idx_t, data_t> & ss_y, double & r_nrm2, double & b_nrm2)
{
    b_nrm2 = vec_dot<idx_t, data_t, double>(ss_b, ss_b);
    b_nrm2 = sqrt(b_nrm2);

    ss_A.Mult(ss_x, ss_y, false);
    vec_add(ss_b, (data_t) -1.0, ss_y, ss_y);
    r_nrm2 = vec_dot<idx_t, data_t, double>(ss_y, ss_y);
    r_nrm2 = sqrt(r_nrm2);
}

void stat_part_times(TEST_RECORD & record)
{
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    int num_procs; MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    MPI_Allreduce(solver->part_times, min_times, NUM_KRYLOV_RECORD, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD);
    MPI_Allreduce(solver->part_times, max_times, NUM_KRYLOV_RECORD, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
    MPI_Allreduce(solver->part_times, avg_times, NUM_KRYLOV_RECORD, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); 
    for (int i = 0; i < NUM_KRYLOV_RECORD; i++)
        avg_times[i] /= num_procs;
    if (my_pid == 0) {
        printf("prec time min/avg/max %.3e %.3e %.3e\n", min_times[PREC], avg_times[PREC], max_times[PREC]);
        printf("oper time min/avg/max %.3e %.3e %.3e\n", min_times[OPER], avg_times[OPER], max_times[OPER]);
        printf("axpy time min/avg/max %.3e %.3e %.3e\n", min_times[AXPY], avg_times[AXPY], max_times[AXPY]);
        printf("dot  tune min/avg/max %.3e %.3e %.3e\n", min_times[DOT ], avg_times[DOT ], max_times[DOT ]);
    }

    record.prec = avg_times[PREC];
    record.oper = avg_times[OPER];
    record.axpy = avg_times[AXPY];
    record.dot  = avg_times[DOT ];
}

void stat_multi_runs(std::vector<TEST_RECORD> & records)
{
    const idx_t test_cnt = records.size();
    std::sort(records.begin(), records.end());
    const TEST_RECORD & best = records[0], & worst = records[test_cnt - 1];
    TEST_RECORD avg;
    for (idx_t i = 0; i < test_cnt; i++) {
        avg.setup += records[i].setup;
        avg.solve += records[i].solve;
        avg.prec  += records[i].prec;
        avg.iter  += records[i].iter;
    }
    avg.setup /= test_cnt;
    avg.solve /= test_cnt;
    avg.prec  /= test_cnt;
    avg.iter  /= test_cnt;
    data_t min_tot = best .setup + best .solve;
    data_t max_tot = worst.setup + worst.solve;
    data_t avg_tot = avg  .setup + avg  .solve;
    printf("\n  Summary of %d tests:\n", test_cnt);
    printf("     Setup time  Solve time  #Iter  Total time\n");
    printf("Min    %.6f    %.6f  %5d  %.6f\n", best .setup, best .solve, best .iter, min_tot);
    printf("Avg    %.6f    %.6f  %5d  %.6f\n", avg  .setup, avg  .solve, avg  .iter, avg_tot);
    printf("Max    %.6f    %.6f  %5d  %.6f\n", worst.setup, worst.solve, worst.iter, max_tot);    
}

void stat_sstructmg_pattern()
{
    MultiGrid<idx_t, data_t, data_t, data_t> * mg_data = dynamic_cast<MultiGrid<idx_t, data_t, data_t, data_t> *>(precond);
    if (mg_data == nullptr) return;

    const idx_t num_levels = mg_data->num_levels();

    for (idx_t ilev = 0; ilev < num_levels; ilev ++) {
        const par_SstructMatrix<idx_t, data_t, data_t>* par_ss_A = mg_data->A_lev(ilev);

        const SStructGrid<idx_t> & ssgrid = par_ss_A->ssgrid;
        const MPI_Comm lev_comm = ssgrid.glb_comm;
        int lev_pid, lev_nprocs;
        MPI_Comm_rank(lev_comm, & lev_pid);
        MPI_Comm_size(lev_comm, & lev_nprocs);

        FILE * fp = nullptr;
        if (lev_pid == 0) { fp = fopen(("sstructmg.L"+std::to_string(ilev)).c_str(), "w"); fclose(fp); }

        for (int p = 0; p < lev_nprocs; p++) {
            if (p == lev_pid) {
                fp = fopen(("sstructmg.L"+std::to_string(ilev)).c_str(), "a");// attach

                const idx_t my_nblks = ssgrid.my_blk_list.size();
                for (idx_t k = 0; k < my_nblks; k++) {
                    const par_structMatrix<idx_t, data_t, data_t> * diag_A = par_ss_A->diag_matrixs[k];
                    const Block<idx_t> & blk = ssgrid.my_blk_list[k];
                    const idx_t my_nboxs = blk.my_box_list.size();

                    for (idx_t b = 0; b < my_nboxs; b++) {
                        const StructCommPackage<idx_t> * str_pkg = diag_A->comm_pkgs[b];
                        const idx_t num_diag = diag_A->local_matrixs[b]->num_diag;

                        const idx_t num_recvs = str_pkg->recv_pids.size();
                        for (idx_t ir = 0; ir < num_recvs; ir++) {
                            const int src_pid = blk.pid_blk2glb->operator[](str_pkg->recv_pids[ir]);// 块内通信域转换为全局通信域
                            int num;
                            MPI_Type_size(str_pkg->recv_subarrs[ir], & num);// return bytes
                            num /= (sizeof(data_t) * num_diag);// 注意要除以对角线数目，因为这个是矩阵的halo区
                            fprintf(fp, "%d %d %d\n", lev_pid, src_pid, num);
                        }
                    }

                    const CCSRMatrix<idx_t, data_t, data_t> * offd_A = par_ss_A->offd_matrixs[k];
                    const UnstructCommUnit<idx_t> * irr_recv_book = offd_A->recv_book;
                    for (idx_t ir = 0; ir < offd_A->num_recvs; ir++) {
                        const int src_pid = irr_recv_book[ir].pid;
                        const int num = irr_recv_book[ir].num_elems;
                        fprintf(fp, "%d %d %d\n", lev_pid, src_pid, num);
                    }
                }
                fclose(fp);
                fflush(stdout);
            }
            MPI_Barrier(lev_comm);
        }// process print sequentially
    }// lev loop
}